package cn.dm.pojo;

import io.lettuce.core.RedisClient;
import io.lettuce.core.ScriptOutputType;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.sync.RedisCommands;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 分布式锁
 * 原链接:https://juejin.cn/post/6844904057920831496
 */
public class Lock {
    // 输出日志
    private Logger logger = LoggerFactory.getLogger(Lock.class);
    // 当前线程的锁集合
    private ThreadLocal<Map<String, Integer>> lockers = new ThreadLocal<>();
    // 当前线程锁的key和value集合
    private ThreadLocal<Map<String, String>> values = new ThreadLocal<>();
    private RedisClient client;
    private RedisCommands<String, String> commands;
    // 两个原子变量，用于存储加锁和解锁脚本的SHA_ID
    private final AtomicReference<String> LUA_SHA_LOCK = new AtomicReference<>();
    private final AtomicReference<String> LUA_SHA_UNLOCK = new AtomicReference<>();

    // 构造函数,形参为RedisClient的实例
    public Lock(RedisClient client) {
        this.client = client;
        StatefulRedisConnection<String, String> connection = client.connect();
        commands = connection.sync();
        // 加锁的lua脚本
        // 键不存在则创建，并设置键的的过期时间为5s, 如果已存在则创建失败，
        final String lockLua = "local result = redis.call('setnx', KEYS[1], ARGV[2]); " +
                "if result == 1 then redis.call('pexpire', KEYS[1], tonumber(ARGV[1])) " +
                "return nil else return redis.call('pttl', KEYS[1]) end";
        // 将脚本缓存到服务器，并保存它的唯一ID到原子变量中
        LUA_SHA_LOCK.compareAndSet(null, commands.scriptLoad(lockLua));
        // 释放锁的lua脚本
        final String unlockLua = "local result = redis.call('get', KEYS[1]);" +
                "if result == ARGV[1] then redis.call('del', KEYS[1]) " +
                "return 1 else return nil end";
        // 将脚本缓存到服务器，并保存它的唯一ID到原子变量中
        LUA_SHA_UNLOCK.compareAndSet(null, commands.scriptLoad(unlockLua));
    }

    /**
     * 尝试加锁
     *
     * @param key 锁的键
     * @return 是否成功获取锁
     */
    private Boolean tryLock(String key) {
        String[] keys = new String[]{key};
        // 失效时间为5秒
        String ttl = "5000";
        String value = getValueByKey(key);
        // 如果没有设置过值
        if (value == null) {
            // 锁的值使用UUID生成随机ID以保证值的唯一性
            value = UUID.randomUUID().toString();
            // 将新生成的值放入集合中
            values.get().put(key, value);
        }
        String[] args = new String[]{ttl, value};
        // 如果创建键成功，则说明加锁成功
        Long result = commands.evalsha(LUA_SHA_LOCK.get(), ScriptOutputType.INTEGER, keys, args);
        if (result == null) {
            return true;
        }
        boolean isLock;
        // 一直阻塞，直到拿到锁
        while (true) {
            try {
                // 这里的实现是有问题的
                // Redisson实现的分布式锁是使用发布订阅实现的(Java的reentrantLock通过队列实现的)
                // 从而可以及时通知其它线程去抢锁
                Thread.sleep(5);
                // 继续尝试获取锁
                isLock = this.tryLock(key);
                if (isLock) {
                    return true;
                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 尝试释放
     *
     * @param key 锁的键
     * @return 是否成功释放锁
     */
    private boolean tryRelease(String key) {
        String[] keys = new String[]{key};
        String[] args = new String[]{getValueByKey(key)};
        // 释放锁
        Long result = commands.evalsha(LUA_SHA_UNLOCK.get(), ScriptOutputType.INTEGER, keys, args);
        return result != null;
    }

    /**
     * 获取当前线程的锁
     *
     * @param key 锁的键
     * @return 当前线程的锁和该锁的重入次数
     */
    private Integer getLockerCnt(String key) {
        // 获取当前线程的锁集合
        Map<String, Integer> map = lockers.get();
        // 如果集合不为空，返回key对应的值
        if (map != null) {
            return map.get(key);
        }
        lockers.set(new HashMap<>(4));
        return null;
    }

    /**
     * 获取锁对应的值
     *
     * @param key 锁的键
     * @return 锁对应的值
     */
    private String getValueByKey(String key) {
        // 获取当前线程的锁和对应值的键值对集合
        Map<String, String> map = values.get();
        // 如果集合不为空，返回key对应的值
        if (map != null) {
            return map.get(key);
        }
        values.set(new HashMap<>(4));
        return null;
    }

    /**
     * 加可重入锁
     *
     * @param key 锁的键
     * @return 是否成功
     */
    public boolean lock(String key) {
        Integer refCnt = getLockerCnt(key);
        if (refCnt != null) {
            // 如果锁已持有，则锁的引用计数加1
            lockers.get().put(key, refCnt + 1);
            return true;
        }
        // 尝试加锁
        boolean ok = this.tryLock(key);
        // 如果加锁失败，则返回
        if (!ok) {
            return false;
        }
        // 加锁成功，引用计数设置为1
        lockers.get().put(key, 1);
        return true;
    }

    /**
     * 释放可重入锁
     *
     * @param key 锁的键
     * @return 是否成功
     */
    public boolean unlock(String key) {
        Integer refCnt = getLockerCnt(key);
        // 当前未持有锁
        if (refCnt == null) {
            return false;
        }
        // 锁的引用数减1
        refCnt--;
        // 引用计数大于0，说明还持有锁
        if (refCnt > 0) {
            lockers.get().put(key, refCnt);
        } else {
            // 否则从锁集合中删除该键，并释放锁
            lockers.get().remove(key);
            return this.tryRelease(key);
        }
        return true;
    }
}

